use chrono::Local;
use clap::Parser;
use ssh2::Session;
use std::io::prelude::*;
use std::{
    fs::File,
    net::TcpStream,
    path::{Path, PathBuf},
    process::exit,
    time::Duration,
};

const SSH_TOTAL_RETRY_COUNT: usize = 10;

#[derive(Debug, Default)]
struct Server {
    pub hostname: String,
    pub port: String,
    pub username: String,
    pub password: String,
}

#[inline]
fn get_filesize(file: &PathBuf) -> u64 {
    std::fs::metadata(file)
        .map(|metadata| metadata.len())
        .unwrap_or(0)
}

#[inline]
fn eta_format(secs: u64) -> String {
    let remaining_seconds = secs % 60;
    let minutes = (secs % 3600) / 60;
    let hours = secs / 3600;
    let s = format!("{:02.0}:{:02.0}:{:02.0}", hours, minutes, remaining_seconds);
    s
}

fn connect_ssh(s: &Server) -> Option<Session> {
    // 新建连接
    let tcp = match TcpStream::connect(format!("{}:{}", s.hostname, s.port)) {
        Ok(tcp) => tcp,
        Err(e) => {
            // 无法链接到对应的端口
            eprintln!(
                "Host: {}:{}, Session create failed, Cause: {}",
                s.hostname, s.port, e
            );
            return None;
        }
    };
    let mut sess = Session::new().unwrap();
    sess.set_tcp_stream(tcp);
    match sess.handshake() {
        Ok(()) => {}
        Err(e) => {
            eprintln!(
                "Host: {}:{}, Server handshake failed, cause: {}",
                s.hostname, s.port, e
            );
            return None;
        }
    }
    let pwd = &s.password.clone();
    match sess.userauth_password(&s.username, pwd) {
        Ok(()) => {}
        Err(e) => {
            eprintln!(
                "Host: {}:{}, Server auth failed, cause: {}",
                s.hostname, s.port, e
            );
            return None;
        }
    }
    println!("Connected to server {}:{}", s.hostname, s.port);
    if let Err(e) = sess.set_banner("monica") {
        eprintln!(
            "Update server {}:{} set_banner error, cause: {}",
            s.hostname, s.port, e
        );
    }
    //
    // sess.set_keepalive(true, SSH_KEEPALIVE_INTERVAL as u32);

    Some(sess)
}

#[inline]
// 向远程服务器发送文件
fn scp_send(sess: Session, s: Server, file: PathBuf, rfile: PathBuf) {
    let remote_file = rfile.to_string_lossy().to_string();
    // let local_file = file.to_string_lossy().to_string();
    let local_file_name = file.file_name().unwrap().to_str().unwrap();
    // info!("xlsx:Line: {:<2} Upload \"{}\"", self.rid, local_file);

    // 获取本地文件的基础信息
    let file_size = get_filesize(&file);
    // Write the file
    // 文件繁忙：
    // called `Result::unwrap()` on an `Err` value: Error { code: Session(-28), msg: "failed to send file" }

    // 避免程序正在使用/运行，先生成临时文件
    let remote_tmp_file_path: &String = &format!("{}.sshscp", remote_file);
    // 临时文件
    let remote_tmp_file = Path::new(remote_tmp_file_path);

    let remote_sftp = sess.sftp().unwrap();
    let remote_dir = rfile.parent().unwrap().iter().as_path();

    if let Err(e) = remote_sftp.readdir(remote_dir) {
        println!("[WARN] sshscp: check dir: {:?}: {e}", remote_dir);
        match remote_sftp.mkdir(remote_dir, 0o755) {
            Ok(()) => println!("[INFO] sshscp: created remote dir: {:?}", remote_dir),
            Err(e) => println!("[ERROR] sshscp: create remote dir failed: {e}")
        }
    }

    // 打开远程文件
    let mut ch;
    loop {
        match sess.scp_send(&remote_tmp_file, 0o755, file_size, None) {
            Ok(c) => {
                ch = c;
                break;
            }
            Err(e) => {
                eprintln!(
                    "sshscp {:10} error, cause: {}, retry again, sleep 3s.",
                    local_file_name, e
                );
                std::thread::sleep(Duration::from_secs(3));
            }
        };
    }

    // 读取本地文件
    let mut f;
    loop {
        match File::open(&file) {
            Ok(file) => {
                f = file;
                break;
            }
            Err(e) => {
                eprintln!(
                    "sshscp {:10} error, cause: {}, retry again, sleep 3s.",
                    local_file_name, e
                );
            }
        }
        std::thread::sleep(Duration::from_secs(3));
    }

    // 开始写入

    // let mut f = File::open(&file).unwrap();
    let mut buf: Vec<u8> = Vec::new();
    f.read_to_end(&mut buf).unwrap();

    // 文件大小
    let file_kb = file_size as f64 / 1024.0;
    let col_size = format!("{:.0}", file_kb).len();
    // 本次已发送字节数
    let mut bytes_send = 0;
    // 全部发送的字节数
    let mut total_bytes_send: usize = 0;
    // 块大小: 
    // 小文件：16k
    // 大文件：1MB
    let chunk_size = if file_kb >= 1024 as f64 { 1024*1024 } else { 16*1024 };

    // 是否上传完成
    // let mut completed = false;

    // 当前的时钟
    let mut clock = Local::now().timestamp();
    // 16KB
    // 1.5M => 1M
    for chunk in buf.chunks(chunk_size) {
        // called `Result::unwrap()` on an `Err` value: Custom { kind: Other, error: "Unable to send channel data" }
        // 网络不稳定会导致报错：error: "Unable to send channel data"
        match ch.write_all(chunk) {
            Ok(()) => {}
            Err(_) => {
                // error!("xlsx:Line: {:<2} Upload [{}/{}] file {} interrupted, Network not available, {}, retry again after 3s.", self.rid, current, counter, local_file_name, e);
                // std::thread::sleep(Duration::from_secs(3));
                // 再次重试还是失败的话，那只能重新上传整个文件
                match ch.write_all(chunk) {
                    Ok(()) => {}
                    Err(e) => {
                        eprintln!(
                            "sshscp {} interrupted, Network not available, {}",
                            local_file_name, e
                        );
                        break;
                    }
                }
            }
        }
        bytes_send += chunk.len();
        total_bytes_send += chunk.len();

        let p = (total_bytes_send as f64 * 100.0 / file_size as f64).floor();
        let total_send_kb = (total_bytes_send as f64 / 1024.0).ceil();

        // 1、如果传输已经到100%，则不用输出如下日志
        if total_bytes_send as u64 >= file_size {
            // 速度：平均一秒的速度
            let speed_kb: f64 = (bytes_send as f64 / 1024.0 / 1.0).ceil();
            println!(
                "TX {} {:>3.0}% {:col_size$.0}KiB / {:.0}KiB {:col_size$.0}KiB/s Done",
                local_file_name, p, total_send_kb, file_kb, speed_kb
            );
            // completed = true;
            break;
        }

        // 当前的时钟：相差多少秒
        let new_clock = Local::now().timestamp();
        let duration_sec = new_clock - clock;
        if duration_sec > 0 {
            // 速度计算
            let mut speed_kb = (bytes_send as f64 / 1024.0 / duration_sec as f64).ceil();
            if speed_kb == 0.0 {
                speed_kb = 1.0;
            }
            // 计算用时
            // 2、传输文件较大，需要拆分成为多次传输，每次打印的信息按秒打印
            let eta_secs = (file_kb - total_send_kb) / speed_kb;
            let eta = eta_format(eta_secs as u64);
            println!(
                "TX {} {:>3.0}% {:col_size$.0}KiB / {:.0}KiB {:col_size$.0}KiB/s {} ETA",
                local_file_name, p, total_send_kb, file_kb, speed_kb, eta
            );
            // 重置计时器
            clock = Local::now().timestamp();
            bytes_send = 0;
        }
    }

    let mut total_try_count = 0;
    if let Err(e) = ch.send_eof() {
        eprintln!("Channel send EOF: cause: {}", e);
        total_try_count = SSH_TOTAL_RETRY_COUNT;
    }
    if let Err(e) = ch.wait_eof() {
        eprintln!("Channel wait EOF error: cause: {}", e);
        total_try_count = SSH_TOTAL_RETRY_COUNT;
    }
    if let Err(e) = ch.close() {
        eprintln!("Channel close error: cause: {}", e);
        total_try_count = SSH_TOTAL_RETRY_COUNT;
    }
    if let Err(e) = ch.wait_close() {
        eprintln!("Channel wait close error: cause: {}", e);
        total_try_count = SSH_TOTAL_RETRY_COUNT;
    }

    let mut try_count = 0;
    while try_count < total_try_count {
        let _sess = connect_ssh(&s);
        if _sess.is_some() {
            scp_send(_sess.unwrap(), s, file.clone(), rfile.clone());
            break;
        }

        eprintln!(
            "Upload file {} interrupted, Network not available, [{}/{}] retry again after 3s.",
            local_file_name,
            try_count + 1,
            total_try_count
        );
        // 休眠3秒，重试
        std::thread::sleep(Duration::from_secs(3));
        try_count += 1;
        if try_count == 10 {
            // 超过重试次数
            // 网络异常
            println!("ERROR: Network not available, exceeding retry attempts, exit now.");
            exit(-1);
        }
    }
}

#[derive(Parser, Clone, Debug)]
struct Opt {
    /// Unix SSH Host
    #[arg(short = 'H', long, env = "SSH_CLI_HOST")]
    host: String,

    /// Unix SSH Username
    #[arg(short, long, env = "SSH_CLI_USER")]
    username: String,

    /// Unix SSH User Password
    #[arg(short, long, env = "SSH_CLI_PASSWORD")]
    password: String,

    /// Unix SSH Port
    #[arg(short = 'P', long, default_value="22")]
    port: String,

    /// Only supports a single command
    #[arg(short = 'B', long, env = "SSH_SCP_BEFORE_COMMAND")]
    before_scp_command: Option<String>,

    #[arg(short = 'A', long, env = "SSH_SCP_AFTER_COMMAND")]
    after_scp_command: Option<String>,

    /// Local file path
    #[arg(short, long)]
    lfile: String,

    /// Remote file path
    #[arg(short, long)]
    rfile: String,
}

fn main() {
    let opt = Opt::parse();

    let mut s = Server::default();
    s.hostname = opt.host.clone();
    s.password = opt.password.clone();
    s.port = opt.port.clone();
    s.username = opt.username.clone();

    let tcp = TcpStream::connect(format!("{}:{}", opt.host, opt.port)).unwrap();
    let mut sess = Session::new().unwrap();
    sess.set_tcp_stream(tcp);
    sess.handshake().unwrap();

    sess.userauth_password(&opt.username, &opt.password)
        .unwrap();
    assert!(sess.authenticated());

    let mut channel2 = sess.channel_session().unwrap();

    if let Some(cmd) = opt.before_scp_command {
        let mut channel = sess.channel_session().unwrap();
        channel.exec(&cmd).unwrap();

        let mut stdout = String::new();
        channel.read_to_string(&mut stdout).unwrap();
        let mut stderr: String = String::new();
        channel.stderr().read_to_string(&mut stderr).unwrap();

        channel.wait_close().unwrap();
        let status = channel.exit_status().unwrap();

        println!("before.info1: cmd: `{}`", &cmd);
        println!("before.info2: stdout: {}", stdout.trim_end_matches("\n"));
        println!("before.info3: stderr: {}", stderr);
        println!("before.info4: exit code: {}", status);

        if status != 0 {
            exit(status);
        }
    }

    scp_send(sess, s, opt.lfile.into(), opt.rfile.into());

    if let Some(cmd) = opt.after_scp_command {
        channel2.exec(&cmd).unwrap();

        let mut stdout = String::new();
        channel2.read_to_string(&mut stdout).unwrap();
        let mut stderr: String = String::new();
        channel2.stderr().read_to_string(&mut stderr).unwrap();

        channel2.wait_close().unwrap();
        let status = channel2.exit_status().unwrap();

        println!("after.info1: cmd: `{}`", &cmd);
        println!("after.info2: stdout: {}", stdout.trim_end_matches("\n"));
        println!("after.info3: stderr: {}", stderr);
        println!("after.info4: exit code: {}", status);

        if status != 0 {
            exit(status);
        }
    }
}
